"""
会话持久化管理模块

负责管理会话的保存、加载和恢复功能。
"""

import json
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any

from rich.console import Console

console = Console()


class SessionManager:
    """会话持久化管理类"""
    
    def __init__(self, project_dir: Optional[Path] = None):
        """
        初始化会话管理器
        
        Args:
            project_dir: 项目目录路径，如果为None则使用当前目录
        """
        self.project_dir = project_dir or Path.cwd()
        self.minicoder_dir = self.project_dir / ".minicoder"
        self.sessions_dir = self.minicoder_dir / "sessions"
        self.checkpoints_dir = self.minicoder_dir / "checkpoints"
        self.generated_dir = self.minicoder_dir / "generated"
        
        # 确保目录存在
        self._ensure_directories()
        
        # 加载项目配置
        self.project_config = self._load_project_config()
        
        # 当前会话信息
        self.current_session_id: Optional[str] = None
        self.session_data: Dict[str, Any] = {}
    
    def _ensure_directories(self) -> None:
        """确保必要的目录存在"""
        for directory in [self.minicoder_dir, self.sessions_dir, self.checkpoints_dir, self.generated_dir]:
            directory.mkdir(parents=True, exist_ok=True)
    
    def _load_project_config(self) -> Dict[str, Any]:
        """加载项目配置"""
        config_file = self.minicoder_dir / "config.yaml"
        if config_file.exists():
            try:
                import yaml
                with open(config_file, 'r', encoding='utf-8') as f:
                    return yaml.safe_load(f)
            except Exception as e:
                console.print(f"[yellow]警告: 无法加载项目配置: {e}[/yellow]")
        
        # 返回默认配置
        return {
            "session": {
                "auto_save": True,
                "max_history": 50,
                "checkpoint_interval": 10
            },
            "code_generation": {
                "auto_save_code": True,
                "code_output_dir": ".minicoder/generated",
                "backup_original": True
            }
        }
    
    def is_project_initialized(self) -> bool:
        """检查项目是否已初始化"""
        return self.minicoder_dir.exists() and (self.minicoder_dir / "config.yaml").exists()
    
    def create_session(self, session_name: Optional[str] = None) -> str:
        """
        创建新会话
        
        Args:
            session_name: 会话名称，如果为None则自动生成
            
        Returns:
            会话ID
        """
        if session_name is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            session_name = f"session_{timestamp}"
        
        session_id = session_name
        session_file = self.sessions_dir / f"{session_id}.json"
        
        # 初始化会话数据
        self.session_data = {
            "id": session_id,
            "name": session_name,
            "created": datetime.now().isoformat(),
            "last_updated": datetime.now().isoformat(),
            "history": [],
            "context_files": [],
            "memory": {
                "entries": [],
                "settings": {
                    "max_entries": 100,
                    "auto_cleanup": True,
                    "categories": ["项目信息", "用户偏好", "技术要点", "重要提醒", "代码片段"]
                }
            },
            "config": {
                "provider": None,
                "model": None,
                "stream": False
            },
            "metadata": {
                "total_messages": 0,
                "total_tokens": 0,
                "project_dir": str(self.project_dir)
            }
        }
        
        self.current_session_id = session_id
        self._save_session()
        
        console.print(f"[green]✅ 创建新会话: {session_name}[/green]")
        return session_id
    
    def load_session(self, session_id: str) -> bool:
        """
        加载指定会话
        
        Args:
            session_id: 会话ID
            
        Returns:
            是否加载成功
        """
        session_file = self.sessions_dir / f"{session_id}.json"
        
        if not session_file.exists():
            console.print(f"[red]❌ 会话不存在: {session_id}[/red]")
            return False
        
        try:
            with open(session_file, 'r', encoding='utf-8') as f:
                self.session_data = json.load(f)
            
            self.current_session_id = session_id
            console.print(f"[green]✅ 加载会话: {self.session_data.get('name', session_id)}[/green]")
            return True
            
        except Exception as e:
            console.print(f"[red]❌ 加载会话失败: {e}[/red]")
            return False
    
    def save_session(self) -> None:
        """保存当前会话"""
        if self.current_session_id and self.session_data:
            self.session_data["last_updated"] = datetime.now().isoformat()
            self._save_session()
    
    def _save_session(self) -> None:
        """内部保存会话方法"""
        if not self.current_session_id:
            return
        
        session_file = self.sessions_dir / f"{self.current_session_id}.json"
        
        try:
            with open(session_file, 'w', encoding='utf-8') as f:
                json.dump(self.session_data, f, ensure_ascii=False, indent=2)
        except Exception as e:
            console.print(f"[red]❌ 保存会话失败: {e}[/red]")
    
    def add_message(self, prompt: str, response: str, metadata: Optional[Dict] = None) -> None:
        """
        添加消息到会话历史
        
        Args:
            prompt: 用户提示词
            response: AI响应
            metadata: 额外元数据
        """
        if not self.current_session_id:
            return
        
        message = {
            "timestamp": datetime.now().isoformat(),
            "prompt": prompt,
            "response": response,
            "metadata": metadata or {}
        }
        
        self.session_data["history"].append(message)
        self.session_data["metadata"]["total_messages"] += 1
        
        # 限制历史记录数量
        max_history = self.project_config.get("session", {}).get("max_history", 50)
        if len(self.session_data["history"]) > max_history:
            self.session_data["history"] = self.session_data["history"][-max_history:]
        
        # 自动保存
        if self.project_config.get("session", {}).get("auto_save", True):
            self.save_session()
    
    def get_history(self) -> List[Dict[str, Any]]:
        """获取会话历史"""
        return self.session_data.get("history", [])
    
    def add_context_file(self, file_path: str, content: str) -> None:
        """
        添加上下文文件
        
        Args:
            file_path: 文件路径
            content: 文件内容
        """
        if not self.current_session_id:
            return
        
        context_files = self.session_data.get("context_files", [])
        
        # 检查是否已存在
        for ctx_file in context_files:
            if ctx_file["path"] == file_path:
                ctx_file["content"] = content
                ctx_file["last_updated"] = datetime.now().isoformat()
                break
        else:
            # 添加新文件
            context_files.append({
                "path": file_path,
                "content": content,
                "added": datetime.now().isoformat(),
                "last_updated": datetime.now().isoformat()
            })
        
        self.session_data["context_files"] = context_files
        self.save_session()
    
    def get_context_files(self) -> List[Dict[str, Any]]:
        """获取上下文文件列表"""
        return self.session_data.get("context_files", [])
    
    def list_sessions(self) -> List[Dict[str, Any]]:
        """列出所有会话"""
        sessions = []
        
        for session_file in self.sessions_dir.glob("*.json"):
            try:
                with open(session_file, 'r', encoding='utf-8') as f:
                    session_data = json.load(f)
                
                sessions.append({
                    "id": session_data.get("id", session_file.stem),
                    "name": session_data.get("name", session_file.stem),
                    "created": session_data.get("created"),
                    "last_updated": session_data.get("last_updated"),
                    "message_count": len(session_data.get("history", [])),
                    "file": session_file
                })
            except Exception as e:
                console.print(f"[yellow]警告: 无法读取会话文件 {session_file}: {e}[/yellow]")
        
        # 按最后更新时间排序
        sessions.sort(key=lambda x: x.get("last_updated", ""), reverse=True)
        return sessions
    
    def create_checkpoint(self, checkpoint_name: Optional[str] = None) -> str:
        """
        创建检查点
        
        Args:
            checkpoint_name: 检查点名称
            
        Returns:
            检查点ID
        """
        if not self.current_session_id:
            raise ValueError("没有活动会话")
        
        if checkpoint_name is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            checkpoint_name = f"checkpoint_{timestamp}"
        
        checkpoint_id = f"{self.current_session_id}_{checkpoint_name}"
        checkpoint_file = self.checkpoints_dir / f"{checkpoint_id}.json"
        
        checkpoint_data = {
            "id": checkpoint_id,
            "name": checkpoint_name,
            "session_id": self.current_session_id,
            "created": datetime.now().isoformat(),
            "session_data": self.session_data.copy()
        }
        
        try:
            with open(checkpoint_file, 'w', encoding='utf-8') as f:
                json.dump(checkpoint_data, f, ensure_ascii=False, indent=2)
            
            console.print(f"[green]✅ 创建检查点: {checkpoint_name}[/green]")
            return checkpoint_id
            
        except Exception as e:
            console.print(f"[red]❌ 创建检查点失败: {e}[/red]")
            raise
    
    def restore_checkpoint(self, checkpoint_id: str) -> bool:
        """
        恢复检查点
        
        Args:
            checkpoint_id: 检查点ID
            
        Returns:
            是否恢复成功
        """
        checkpoint_file = self.checkpoints_dir / f"{checkpoint_id}.json"
        
        if not checkpoint_file.exists():
            console.print(f"[red]❌ 检查点不存在: {checkpoint_id}[/red]")
            return False
        
        try:
            with open(checkpoint_file, 'r', encoding='utf-8') as f:
                checkpoint_data = json.load(f)
            
            self.session_data = checkpoint_data["session_data"]
            self.current_session_id = checkpoint_data["session_id"]
            
            console.print(f"[green]✅ 恢复检查点: {checkpoint_data.get('name', checkpoint_id)}[/green]")
            return True
            
        except Exception as e:
            console.print(f"[red]❌ 恢复检查点失败: {e}[/red]")
            return False
    
    def auto_save_code(self, response: str, prompt: str) -> List[str]:
        """
        智能保存AI响应中的代码到项目根目录
        只在用户明确要求生成代码时才保存
        
        Args:
            response: AI响应文本
            prompt: 原始提示词
            
        Returns:
            保存的文件路径列表
        """
        if not self.project_config.get("code_generation", {}).get("auto_save_code", True):
            return []
        
        # 检查用户是否明确要求生成代码
        if not self._should_auto_save_code(prompt, response):
            return []
        
        import re
        saved_files = []
        
        # 检测代码块模式：```language\ncode\n```
        code_block_pattern = r'```(\w+)?\n(.*?)```'
        matches = re.findall(code_block_pattern, response, re.DOTALL)
        
        for i, (language, code) in enumerate(matches):
            if not code.strip():
                continue
            
            # 确定文件扩展名
            extension = self._get_extension_for_language(language)
            
            # 智能生成文件名
            filename = self._generate_smart_filename(prompt, language, extension, i)
            
            # 保存到项目根目录
            output_file = self.project_dir / filename
            try:
                # 确保目录存在
                output_file.parent.mkdir(parents=True, exist_ok=True)
                
                with open(output_file, 'w', encoding='utf-8') as f:
                    f.write(code.strip())
                
                # 如果是可执行文件，设置执行权限
                if extension in ['.py', '.sh', '.js']:
                    os.chmod(output_file, 0o755)
                
                saved_files.append(str(output_file))
                console.print(f"[green]💾 保存项目文件: {filename}[/green]")
                
            except Exception as e:
                console.print(f"[red]❌ 保存代码失败: {e}[/red]")
        
        return saved_files
    
    def _should_auto_save_code(self, prompt: str, response: str) -> bool:
        """
        判断是否应该自动保存代码
        
        Args:
            prompt: 用户提示词
            response: AI响应
            
        Returns:
            是否应该保存代码
        """
        prompt_lower = prompt.lower()
        response_lower = response.lower()
        
        # 明确的代码生成请求关键词
        code_generation_keywords = [
            '生成代码', '写代码', '创建文件', '创建脚本', '编写程序',
            'generate code', 'write code', 'create file', 'create script',
            '编写', '生成', '创建', '实现', '开发', '编程',
            'write', 'generate', 'create', 'implement', 'develop', 'program'
        ]
        
        # 明确的非代码生成请求关键词
        non_code_keywords = [
            '分析', '解释', '说明', '介绍', '查看', '搜索', '查找',
            'analyze', 'explain', 'describe', 'introduce', 'view', 'search', 'find',
            '是什么', '做什么', '如何', '为什么', '怎么样',
            'what is', 'what does', 'how', 'why', 'how to'
        ]
        
        # 检查是否包含代码生成关键词
        has_code_keywords = any(keyword in prompt_lower for keyword in code_generation_keywords)
        
        # 检查是否包含非代码生成关键词
        has_non_code_keywords = any(keyword in prompt_lower for keyword in non_code_keywords)
        
        # 如果明确是非代码请求，不保存
        if has_non_code_keywords and not has_code_keywords:
            return False
        
        # 如果明确是代码生成请求，保存
        if has_code_keywords:
            return True
        
        # 检查响应中是否包含实际的代码块
        import re
        code_block_pattern = r'```(\w+)?\n(.*?)```'
        matches = re.findall(code_block_pattern, response, re.DOTALL)
        
        # 如果响应中有代码块，且代码块不是示例或说明性的，则保存
        if matches:
            for language, code in matches:
                code = code.strip()
                # 检查代码块是否足够长且不是简单的示例
                if len(code) > 50 and not self._is_example_code(code):
                    return True
        
        return False
    
    def _is_example_code(self, code: str) -> bool:
        """
        判断代码是否是示例代码
        
        Args:
            code: 代码内容
            
        Returns:
            是否是示例代码
        """
        # 示例代码的特征
        example_indicators = [
            '# 示例', '# example', '# 演示', '# demo',
            'print("hello world")', 'print("Hello, World!")',
            'console.log("hello")', 'System.out.println("hello")',
            'echo "hello"', 'puts "hello"'
        ]
        
        code_lower = code.lower()
        return any(indicator.lower() in code_lower for indicator in example_indicators)
    
    def _generate_smart_filename(self, prompt: str, language: str, extension: str, index: int) -> str:
        """
        根据提示词智能生成文件名
        
        Args:
            prompt: 用户提示词
            language: 编程语言
            extension: 文件扩展名
            index: 文件索引
            
        Returns:
            生成的文件名
        """
        # 常见项目文件名的关键词映射
        project_keywords = {
            'main': ['main', '主', '主要'],
            'app': ['app', '应用', 'application'],
            'server': ['server', '服务', '服务器'],
            'client': ['client', '客户端'],
            'index': ['index', '首页', '入口'],
            'config': ['config', '配置', '设置'],
            'utils': ['utils', '工具', 'utility'],
            'test': ['test', '测试'],
            'readme': ['readme', '说明', '文档']
        }
        
        # 根据语言确定默认文件名
        default_names = {
            'python': 'main.py',
            'javascript': 'app.js',
            'java': 'Main.java',
            'html': 'index.html',
            'css': 'style.css',
            'sql': 'database.sql',
            'bash': 'script.sh',
            'shell': 'script.sh'
        }
        
        # 分析提示词，寻找项目相关的关键词
        prompt_lower = prompt.lower()
        
        # 检查是否包含项目关键词
        for keyword, variations in project_keywords.items():
            for variation in variations:
                if variation in prompt_lower:
                    if keyword == 'main':
                        return f"main{extension}"
                    elif keyword == 'app':
                        return f"app{extension}"
                    elif keyword == 'server':
                        return f"server{extension}"
                    elif keyword == 'client':
                        return f"client{extension}"
                    elif keyword == 'index':
                        return f"index{extension}"
                    elif keyword == 'config':
                        return f"config{extension}"
                    elif keyword == 'utils':
                        return f"utils{extension}"
                    elif keyword == 'test':
                        return f"test{extension}"
                    elif keyword == 'readme':
                        return f"README{extension}"
        
        # 如果没有找到关键词，使用默认命名规则
        if index == 0:
            # 第一个文件使用默认名称
            return default_names.get(language, f"main{extension}")
        else:
            # 后续文件使用描述性名称
            return f"file_{index+1}{extension}"
    
    def _get_extension_for_language(self, language: str) -> str:
        """根据语言获取文件扩展名"""
        language_map = {
            "python": ".py",
            "javascript": ".js",
            "typescript": ".ts",
            "java": ".java",
            "cpp": ".cpp",
            "c": ".c",
            "go": ".go",
            "rust": ".rs",
            "php": ".php",
            "ruby": ".rb",
            "swift": ".swift",
            "kotlin": ".kt",
            "scala": ".scala",
            "html": ".html",
            "css": ".css",
            "sql": ".sql",
            "bash": ".sh",
            "shell": ".sh",
            "yaml": ".yaml",
            "yml": ".yml",
            "json": ".json",
            "xml": ".xml",
            "markdown": ".md",
            "md": ".md"
        }
        
        return language_map.get(language.lower(), ".txt")
    
    def get_session_info(self) -> Dict[str, Any]:
        """获取当前会话信息"""
        if not self.current_session_id:
            return {}
        
        return {
            "id": self.current_session_id,
            "name": self.session_data.get("name", ""),
            "created": self.session_data.get("created", ""),
            "last_updated": self.session_data.get("last_updated", ""),
            "message_count": len(self.session_data.get("history", [])),
            "context_files_count": len(self.session_data.get("context_files", [])),
            "memory_count": len(self.session_data.get("memory", {}).get("entries", [])),
            "project_dir": str(self.project_dir)
        }
    
    # ==================== 内存管理方法 ====================
    
    def add_memory(self, content: str, category: str = "重要提醒", tags: Optional[List[str]] = None, importance: str = "medium") -> str:
        """
        添加记忆条目
        
        Args:
            content: 记忆内容
            category: 分类
            tags: 标签列表
            importance: 重要性 (high/medium/low)
            
        Returns:
            记忆条目ID
        """
        if not self.current_session_id:
            raise ValueError("没有活动会话")
        
        # 生成唯一ID
        memory_id = f"mem_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{len(self.session_data.get('memory', {}).get('entries', []))}"
        
        memory_entry = {
            "id": memory_id,
            "content": content.strip(),
            "category": category,
            "tags": tags or [],
            "created": datetime.now().isoformat(),
            "last_accessed": datetime.now().isoformat(),
            "access_count": 0,
            "importance": importance
        }
        
        # 确保内存结构存在
        if "memory" not in self.session_data:
            self.session_data["memory"] = {"entries": [], "settings": {}}
        
        self.session_data["memory"]["entries"].append(memory_entry)
        
        # 检查是否需要清理
        self._cleanup_memory_if_needed()
        
        self.save_session()
        return memory_id
    
    def get_memory_entries(self, category: Optional[str] = None, tags: Optional[List[str]] = None, importance: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        获取记忆条目
        
        Args:
            category: 分类筛选
            tags: 标签筛选
            importance: 重要性筛选
            
        Returns:
            记忆条目列表
        """
        if not self.current_session_id:
            return []
        
        entries = self.session_data.get("memory", {}).get("entries", [])
        
        # 应用筛选条件
        filtered_entries = []
        for entry in entries:
            # 分类筛选
            if category and entry.get("category") != category:
                continue
            
            # 重要性筛选
            if importance and entry.get("importance") != importance:
                continue
            
            # 标签筛选
            if tags:
                entry_tags = entry.get("tags", [])
                if not any(tag in entry_tags for tag in tags):
                    continue
            
            filtered_entries.append(entry)
        
        # 按最后访问时间排序
        filtered_entries.sort(key=lambda x: x.get("last_accessed", ""), reverse=True)
        return filtered_entries
    
    def update_memory_access(self, memory_id: str) -> None:
        """更新记忆条目的访问信息"""
        if not self.current_session_id:
            return
        
        entries = self.session_data.get("memory", {}).get("entries", [])
        for entry in entries:
            if entry.get("id") == memory_id:
                entry["last_accessed"] = datetime.now().isoformat()
                entry["access_count"] = entry.get("access_count", 0) + 1
                break
        
        self.save_session()
    
    def delete_memory(self, memory_id: str) -> bool:
        """
        删除记忆条目
        
        Args:
            memory_id: 记忆条目ID
            
        Returns:
            是否删除成功
        """
        if not self.current_session_id:
            return False
        
        entries = self.session_data.get("memory", {}).get("entries", [])
        for i, entry in enumerate(entries):
            if entry.get("id") == memory_id:
                del entries[i]
                self.save_session()
                return True
        
        return False
    
    def refresh_memory(self) -> Dict[str, int]:
        """
        刷新内存，清理过期和低价值条目
        
        Returns:
            清理统计信息
        """
        if not self.current_session_id:
            return {"deleted": 0, "merged": 0}
        
        entries = self.session_data.get("memory", {}).get("entries", [])
        settings = self.session_data.get("memory", {}).get("settings", {})
        
        deleted_count = 0
        merged_count = 0
        
        # 1. 删除低价值条目（访问次数少且创建时间久远）
        current_time = datetime.now()
        entries_to_keep = []
        
        for entry in entries:
            created_time = datetime.fromisoformat(entry.get("created", ""))
            access_count = entry.get("access_count", 0)
            importance = entry.get("importance", "medium")
            
            # 保留条件：高重要性、最近访问、或访问次数多
            days_old = (current_time - created_time).days
            should_keep = (
                importance == "high" or
                access_count > 2 or
                days_old < 7
            )
            
            if should_keep:
                entries_to_keep.append(entry)
            else:
                deleted_count += 1
        
        # 2. 合并相似条目
        merged_entries = []
        for entry in entries_to_keep:
            content = entry.get("content", "").lower()
            category = entry.get("category", "")
            
            # 查找相似的条目
            merged = False
            for merged_entry in merged_entries:
                if (merged_entry.get("category") == category and
                    self._calculate_similarity(content, merged_entry.get("content", "").lower()) > 0.7):
                    # 合并内容
                    merged_entry["content"] = f"{merged_entry['content']}\n\n{entry['content']}"
                    merged_entry["tags"] = list(set(merged_entry.get("tags", []) + entry.get("tags", [])))
                    merged_entry["access_count"] = max(merged_entry.get("access_count", 0), entry.get("access_count", 0))
                    merged = True
                    merged_count += 1
                    break
            
            if not merged:
                merged_entries.append(entry)
        
        # 更新内存
        self.session_data["memory"]["entries"] = merged_entries
        self.save_session()
        
        return {"deleted": deleted_count, "merged": merged_count}
    
    def _cleanup_memory_if_needed(self) -> None:
        """如果需要，清理内存条目"""
        settings = self.session_data.get("memory", {}).get("settings", {})
        max_entries = settings.get("max_entries", 100)
        auto_cleanup = settings.get("auto_cleanup", True)
        
        if not auto_cleanup:
            return
        
        entries = self.session_data.get("memory", {}).get("entries", [])
        if len(entries) <= max_entries:
            return
        
        # 按重要性、访问次数、创建时间排序，保留最重要的条目
        entries.sort(key=lambda x: (
            {"high": 3, "medium": 2, "low": 1}.get(x.get("importance", "medium"), 2),
            x.get("access_count", 0),
            x.get("created", "")
        ), reverse=True)
        
        # 保留前 max_entries 个条目
        self.session_data["memory"]["entries"] = entries[:max_entries]
    
    def _calculate_similarity(self, text1: str, text2: str) -> float:
        """计算两个文本的相似度（简单的词汇重叠度）"""
        if not text1 or not text2:
            return 0.0
        
        words1 = set(text1.split())
        words2 = set(text2.split())
        
        if not words1 or not words2:
            return 0.0
        
        intersection = words1.intersection(words2)
        union = words1.union(words2)
        
        return len(intersection) / len(union) if union else 0.0
    
    def get_memory_categories(self) -> List[str]:
        """获取所有记忆分类"""
        if not self.current_session_id:
            return []
        
        settings = self.session_data.get("memory", {}).get("settings", {})
        return settings.get("categories", ["项目信息", "用户偏好", "技术要点", "重要提醒", "代码片段"])
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """获取内存统计信息"""
        if not self.current_session_id:
            return {}
        
        entries = self.session_data.get("memory", {}).get("entries", [])
        
        # 统计各分类的条目数量
        category_stats = {}
        importance_stats = {"high": 0, "medium": 0, "low": 0}
        total_access_count = 0
        
        for entry in entries:
            category = entry.get("category", "未分类")
            importance = entry.get("importance", "medium")
            access_count = entry.get("access_count", 0)
            
            category_stats[category] = category_stats.get(category, 0) + 1
            importance_stats[importance] = importance_stats.get(importance, 0) + 1
            total_access_count += access_count
        
        return {
            "total_entries": len(entries),
            "category_stats": category_stats,
            "importance_stats": importance_stats,
            "total_access_count": total_access_count,
            "average_access_count": total_access_count / len(entries) if entries else 0
        }
